import gzip
import websockets
import json
import asyncio
import threading
import uuid
from flask import Flask, request, Response, jsonify
from enum import IntEnum
from queue import Queue, Empty

from azure_util import speech_to_text
from latex2chinese import LatexToChinese
import signature as sig
from config import API


# 定义事件枚举
class Event(IntEnum):
    NoneEvent = 0
    StartConnection = 1
    FinishConnection = 2
    ConnectionStarted = 50
    ConnectionFailed = 51
    ConnectionFinished = 52
    StartSession = 100
    FinishSession = 102
    SessionStarted = 150
    SessionFinished = 152
    SessionFailed = 153
    TaskRequest = 200
    TTSSentenceStart = 350
    TTSSentenceEnd = 351
    TTSResponse = 352


# 事件名称映射
event_name = {
    0: "None",
    1: "StartConnection",
    2: "FinishConnection",
    50: "ConnectionStarted",
    51: "ConnectionFailed",
    52: "ConnectionFinished",
    100: "StartSession",
    102: "FinishSession",
    150: "SessionStarted",
    152: "SessionFinished",
    153: "SessionFailed",
    200: "TaskRequest",
    350: "TTSSentenceStart",
    351: "TTSSentenceEnd",
    352: "TTSResponse",
}

# 常量定义
HOST = "openspeech.bytedance.com"
endpoint = "v3/tts/bidirection"
resource_id = "volc.service_type.10029"
app_key = "8333084809"
access_key = "fvS6p_47u0WPDtXhfU7DUweXgIGa0deq"

SPEAKER_MEI_LI_NV_YOU = "zh_female_shuangkuaisisi_moon_bigtts"
namespece = "BidirectionalTTS"
sample_rate = 16000
speech_rate = 0

app = Flask(__name__)


def parse_response(response_bytes):
    print("开始解析数据")
    # 解析头部信息
    protocol_version = (response_bytes[0] >> 4) & 0x0F  # 版本
    header_size = response_bytes[0] & 0x0F  # 头部大小
    message_type = (response_bytes[1] >> 4) & 0x0F  # 消息类型
    message_flags = response_bytes[1] & 0x0F  # 消息标志
    serialization_method = (response_bytes[2] >> 4) & 0x0F  # 序列化方法
    compression_method = response_bytes[2] & 0x0F  # 压缩方法
    # 解析事件编号（4字节）
    event_type = int.from_bytes(response_bytes[4:8], byteorder='big')

    # 解析连接ID长度（4字节）
    connection_id_len = int.from_bytes(response_bytes[8:12], byteorder='big')
    connection_id = response_bytes[12:12 + connection_id_len].decode('utf-8')

    # 解析负载长度（4字节）
    payload_len_start = 12 + connection_id_len
    payload_len = int.from_bytes(response_bytes[payload_len_start:payload_len_start + 4], byteorder='big')

    # 解析负载（JSON）
    payload_start = payload_len_start + 4
    payload_bytes = response_bytes[payload_start:payload_start + payload_len]
    succeed_payload = {"message": "success", "status_code": 0, "connection_id": connection_id, "event_type": event_type}
    if event_type == Event.ConnectionFinished:
        # 解析JSON格式的错误信息
        payload = payload_bytes.decode('utf-8')
        response_meta = json.loads(payload)
        print(connection_id)
        print(response_meta)
        status_code = response_meta.get("status_code", 0)
        message = response_meta.get("message", "")
        return {
            "connection_id": connection_id,
            "status_code": status_code,
            "message": message,
            "event_type": event_type,
        }
    elif event_type in [Event.ConnectionStarted, Event.SessionStarted, Event.SessionFinished, Event.ConnectionFinished]:
        return succeed_payload
    elif event_type == Event.TTSResponse:
        # 音频字节流
        return payload_bytes
    elif event_type in [Event.ConnectionFailed, Event.SessionFailed]:
        # 解析JSON格式的错误信息
        payload = payload_bytes.decode('utf-8')
        response_meta = json.loads(payload)
        print(response_meta)
        status_code = response_meta.get("status_code", 0)
        message = response_meta.get("message", "")
        return {
            "connection_id": connection_id,
            "status_code": status_code,
            "message": message,
            "event_type": event_type,
        }
    else:
        return {"status": "unknown", "connection_id": connection_id, "event_type": event_type}


def generate_request_data(event_type, session_id=None, payload={}):
    # 版本: 0b0001
    protocol_version = 0b0001
    # 头部大小: 0b0001 (4 bytes)
    header_size = 0b0001
    # 消息类型: 0b0001 (完整客户端请求)
    message_type = 0b0001
    # 消息类型特定标志: 0b0100 (with event number)
    message_type_flags = 0b0100
    # 序列化方法: 0b0001 (JSON)
    serialization_method = 0b0001
    # 压缩方法: 0b0000 (无压缩)
    compression_method = 0b0000
    # 保留位: 0b00000000
    reserved = 0b00000000
    # 构建新的 header
    header = bytearray(4)  # 4 字节的 header

    # 将各部分填充到 header 中
    header[0] = (protocol_version << 4) | header_size  # 0 ~ 1 字节
    header[1] = (message_type << 4) | message_type_flags  # 2 ~ 3 字节
    header[2] = (serialization_method << 4) | compression_method  # 4 ~ 5 字节
    header[3] = reserved  # 6 ~ 7 字节
    full_client_request = bytearray(header)

    # header
    full_client_request.extend((event_type).to_bytes(4, 'big'))  # event type(4 bytes)

    # session_id
    if session_id is not None:
        full_client_request.extend((len(session_id)).to_bytes(4, 'big'))  # session id length(4 bytes)
        full_client_request.extend(session_id.encode())  # session id

    # payload
    payload['event'] = event_type.value
    payload_bytes = json.dumps(payload).encode()
    full_client_request.extend((len(payload_bytes)).to_bytes(4, 'big'))  # payload size(4 bytes)
    full_client_request.extend(payload_bytes)  # payload

    return full_client_request


async def start_connection(websocket):
    start_connection_data = generate_request_data(Event.StartConnection)
    await websocket.send(start_connection_data)

    response = await websocket.recv()
    response = parse_response(response)
    print("start_connection response:", response)


async def start_session(websocket, session_id):
    request_json = {
        "namespace": namespece,
        "req_params": {
            "speaker": SPEAKER_MEI_LI_NV_YOU,
            "audio_params": {
                "format": "mp3",
                "sample_rate": sample_rate,
                "speech_rate": speech_rate
            }
        }
    }

    start_session_data = generate_request_data(Event.StartSession, session_id=session_id, payload=request_json)
    await websocket.send(start_session_data)
    response = await websocket.recv()
    response = parse_response(response)
    print("start_session response:", response)


async def send_tts_message(websocket, session_id, text):
    request_json = {
        "namespace": namespece,
        "req_params": {
            "speaker": SPEAKER_MEI_LI_NV_YOU,
            "text": text,
            "audio_params": {
                "format": "mp3",
                "sample_rate": sample_rate,
                "speech_rate": speech_rate
            }
        }
    }
    tts_request_data = generate_request_data(Event.TaskRequest, session_id=session_id, payload=request_json)
    print(f"send_tts_message: {request_json}")
    await websocket.send(tts_request_data)


async def finish_session(websocket, session_id):
    finish_session_data = generate_request_data(Event.FinishSession, session_id=session_id)
    await websocket.send(finish_session_data)

    response = await websocket.recv()
    response = parse_response(response)
    print("finish_session response:", response)


async def finish_connection(websocket):
    finish_connection_data = generate_request_data(Event.FinishConnection)
    await websocket.send(finish_connection_data)

    response = await websocket.recv()
    response = parse_response(response)
    print("finish_connection response:", response)


async def tts_request(text, session_id):
    headers = {
        "X-Api-App-Key": app_key,
        "X-Api-Access-Key": access_key,
        "X-Api-Resource-Id": resource_id,
        "X-Api-Request-Id": session_id,
        "X-Api-Connect-Id": session_id
    }
    # WebSocket 连接的处理
    async with websockets.connect(f"wss://{HOST}/api/v3/tts/bidirection", extra_headers=headers) as websocket:
        await start_connection(websocket)
        await start_session(websocket, session_id)
        # 发送 TTS 请求
        await send_tts_message(websocket, session_id, text)

        await finish_session(websocket, session_id)
        # 接收音频数据
        print("开始接收音频数据...")
        while True:
            response = await websocket.recv()
            # print("收到响应")
            tts_response = parse_response(response)
            # 判定是否是字典
            if isinstance(tts_response, dict):
                event_type = tts_response.get("event_type")
                if event_type == Event.TTSSentenceStart:
                    print("sentence start")
                    pass
                elif event_type == Event.TTSSentenceEnd:
                    print("sentence end")
                    # break
                elif event_type == Event.SessionFinished:
                    print("session finished")
                    break
            else:
                # 纯音频流
                # print("audio data")
                yield tts_response

        await finish_connection(websocket)


# 异步函数，运行 TTS 并将数据放入队列
async def async_tts_request(text, session_id, q):
    try:
        async for audio_chunk in tts_request(text, session_id):
            q.put(audio_chunk)
    except Exception as e:
        print(f"Error in async_tts_request: {e}")
    finally:
        q.put(None)  # 放入结束信号


@app.route('/tts', methods=['POST'])
async def tts_api():
    try:
        # 提取 headers,    # 提取 body
        headers = dict(request.headers)
        body = request.get_json()
        print(type(body))
        if not await sig.verifyRequestSignature(body,headers,API['REQUEST_SIGNATURE_KEY']):
            return jsonify({"code": 401, "error": "signature is invalid !"}), 401
        else:
            print('signature verified ok!')


        data = request.json
        text = data.get("text")
        if not text:
            return jsonify({"code": 400, "message": "text is required"}), 400

        session_id = str(uuid.uuid4())
        q = Queue()
        
        newText = LatexToChinese.process_latex(text)
        print(f"newText: {newText}")

        # 在单独的线程中运行异步 TTS 函数
        def run_async_tts():
            asyncio.run(async_tts_request(newText, session_id, q))

        # 启动异步 TTS 处理的新线程
        thread = threading.Thread(target=run_async_tts)
        thread.start()

        # 生成器函数，从队列中读取音频数据并生成响应
        def generate_audio():
            while True:
                try:
                    data = q.get(timeout=5)  # 设置超时避免无限阻塞
                    if data is None:  # 检查是否收到结束信号
                        print('no data left, end connection')
                        break
                    yield data
                except Empty:
                    break

        # 返回流式响应
        
        return Response(generate_audio(), mimetype='audio/mpeg')
    except Exception as e:
        print(f"Error in tts_api: {e}")
        return jsonify({"error": "tts api error"}), 500


@app.route('/azure', methods=['POST'])
async def azure_api():
    try:
        # 提取 headers,    # 提取 body
        headers = dict(request.headers)

        if 'file' not in request.files:
            return jsonify({"code": 400, "error": "No file part in the request"}), 400

        file = request.files['file']
        if file.filename == '':
            return jsonify({"code": 400, "error": "No selected file"}), 400
        
        reqBody = {'name': file.filename}
        
        if not await sig.verifyRequestSignature(reqBody, headers, API['REQUEST_SIGNATURE_KEY']):
            return jsonify({"code": 401, "error": "signature is invalid !"}), 401
        else:
            print('signature verified ok!')

        if file:
            # 将文件内容读取为字节流
            audio_data = file.read()  # 读取文件内容为字节
            # 在这里将字节数据传递给 speech_to_text 函数
            text = await speech_to_text(audio_data)
            print(text)
            return jsonify({"code": 0, "message": "success", "data": text}), 200
    except Exception as e:
        print(f"Error in azure_api: {e}")
        return jsonify({"code": 500, "message": "azure api error"}), 500


@app.get('/')
def index():
    return "Hello World!"


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5500)
